from os import path, makedirs
from dolfin import File, Expression, interpolate
from numpy import amax
import sys

class SolverBase:
    'Base class for all solvers.'
    def __init__(self,p=1):
        self.solverDir = self.__class__.__name__.lower()
        self.saveDir = None             # results directory                   
        
        self.p = p                      # polynomial degree

        self.phiOut = 'phi.pvd'         # solution file
        self.v = None                   # velocity expression
        self.CFL = 0.125                 # time stepping constant
        self.sf = 10                    # save frequency for solution
        self.pgb = True                 # progress bar
    
    def _get_velocity(self,problem,space,t):
        '''Return interpolated expression for velocity at time t.'''
        if self.v == None:
             vx = problem.vStrings[0]
             vy = problem.vStrings[1]
             self.v = Expression((vx,vy),t = t)

        self.v.t = t
        return interpolate(self.v,space)

    def _get_time_step(self,problem,space,t):
        '''Return time step based on CFL.'''
        v = self._get_velocity(problem,space,t)
        return self.CFL*space.mesh().hmin()/amax(v.vector().array())
    
    def update(self,problem,phi,t,dt,count):
        '''Update method.'''
        if type(self.phiOut) == str:
            self.saveDir = '/'.join([problem.resultsDir,problem.problemDir,self.methodDir,self.solverDir])
            self.phiOut = '/'.join([self.saveDir,'_'.join([str(self.CFL),self.phiOut])])
            self.phiOut = File(self.phiOut)
        
            if not path.exists(self.saveDir):
                makedirs(self.saveDir)
        
        if count % self.sf == 0:   # self.sf increased makes things less verbose
            problem.update(phi,t,dt)
            self.phiOut << phi, count/self.sf
            if self.pgb:
                sys.stdout.write('\r%d%%' % int(100*t/problem.T))
                sys.stdout.flush()

#<----------------------------------------------METHODS TO BE DEFINED IN CHILDREN-------------------------------------------------

    def solve(self,problem):
        '''Solving loop.'''
        raise NotImplementedError('Method solve() must be defined in child classes.')


